import cv2
import numpy as np
from keras.layers import Input
from model import VGG16
import matplotlib.pyplot as plt
import tensorflow as tf

# 加载模型
target_size = (256, 256)
model_input = Input(shape=(target_size[0], target_size[1], 3))
dropout = True
with_CPFE = True
with_CA = True
with_SA = True
model = VGG16(model_input, dropout=dropout, with_CPFE=with_CPFE, with_CA=with_CA, with_SA=with_SA)

# 加载训练好的模型权重
model.load_weights('model/PFA_00050.weights.h5')

# 读取单张照片
image_path = '0001.jpg'  # 替换照片路径
x = cv2.imread(image_path)
x = np.array(x, dtype=np.float32)

# 预处理图像
x = x[..., ::-1]
# Zero-center by mean pixel
x[..., 0] -= 103.939
x[..., 1] -= 116.779
x[..., 2] -= 123.68

# 调整图像大小
x = cv2.resize(x, target_size, interpolation=cv2.INTER_LINEAR)
x = np.expand_dims(x, axis=0)  # 添加批次维度

# 进行预测
prediction = model.predict(x)
prediction = tf.sigmoid(prediction).numpy()  # 应用sigmoid函数
prediction = prediction.reshape(target_size)  # 调整预测结果的形状

# 显示识别前和识别后的图像
plt.figure(figsize=(10, 5))

# 显示原始图像
plt.subplot(1, 2, 1)
plt.imshow(cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB))
plt.title('Original Image')
plt.axis('off')

# 显示识别结果
plt.subplot(1, 2, 2)
plt.imshow(prediction, cmap='gray')
plt.title('Predicted Saliency Map')
plt.axis('off')

plt.show()